const fs = require('fs');
const { getUbyteFileData } = require('./parseMinist');

// 激活函数 sigmoid，用于引入非线性，将输入值映射到 (0, 1) 区间
function sigmoid(x) {
  return 1 / (1 + Math.exp(-x));
}

// sigmoid 函数的导数，在反向传播时用于计算梯度
function sigmoidDerivative(x) {
  return x * (1 - x);
}

// 将一维数组转换为二维数组，方便后续卷积和池化操作
// 参数 data 是一维数组，width 是二维数组每一行的元素个数
function convert1Dto2D(data, width) {
  // 计算二维数组的行数
  const height = data.length / width;
  const result = [];
  // 遍历每一行
  for (let i = 0; i < height; i++) {
    // 从一维数组中截取对应行的数据添加到二维数组中
    result.push(data.slice(i * width, (i + 1) * width));
  }
  return result;
}

// 卷积操作，将输入数据与卷积核进行卷积运算
// 参数 input 是输入的二维数据，kernel 是卷积核，stride 是卷积步长
function convolve(input, kernel, stride = 1) {
  // 获取输入数据的高度和宽度
  const inputHeight = input.length;
  const inputWidth = input[0].length;
  // 获取卷积核的高度和宽度
  const kernelHeight = kernel.length;
  const kernelWidth = kernel[0].length;
  // 计算卷积输出的高度和宽度
  const outputHeight = Math.floor((inputHeight - kernelHeight) / stride) + 1;
  const outputWidth = Math.floor((inputWidth - kernelWidth) / stride) + 1;

  const output = [];
  // 遍历输出的每一行
  for (let y = 0; y < outputHeight; y++) {
    output[y] = [];
    // 遍历输出的每一列
    for (let x = 0; x < outputWidth; x++) {
      let sum = 0;
      // 遍历卷积核的每一行
      for (let ky = 0; ky < kernelHeight; ky++) {
        // 遍历卷积核的每一列
        for (let kx = 0; kx < kernelWidth; kx++) {
          // 计算卷积核元素与输入数据对应位置元素的乘积并累加
          sum += input[y * stride + ky][x * stride + kx] * kernel[ky][kx];
        }
      }
      // 将累加结果存入输出数组
      output[y][x] = sum;
    }
  }
  return output;
}

// 最大池化操作，对输入数据进行下采样，减少数据维度
// 参数 input 是输入的二维数据，poolSize 是池化窗口大小，stride 是池化步长
function maxPooling(input, poolSize = 2, stride = 2) {
  // 获取输入数据的高度和宽度
  const inputHeight = input.length;
  const inputWidth = input[0].length;
  // 计算池化输出的高度和宽度
  const outputHeight = Math.floor((inputHeight - poolSize) / stride) + 1;
  const outputWidth = Math.floor((inputWidth - poolSize) / stride) + 1;

  const output = [];
  // 遍历输出的每一行
  for (let y = 0; y < outputHeight; y++) {
    output[y] = [];
    // 遍历输出的每一列
    for (let x = 0; x < outputWidth; x++) {
      let maxValue = -Infinity;
      // 遍历池化窗口的每一行
      for (let py = 0; py < poolSize; py++) {
        // 遍历池化窗口的每一列
        for (let px = 0; px < poolSize; px++) {
          // 获取池化窗口内的元素值
          const value = input[y * stride + py][x * stride + px];
          // 如果该值大于当前最大值，则更新最大值
          if (value > maxValue) {
            maxValue = value;
          }
        }
      }
      // 将最大值存入输出数组
      output[y][x] = maxValue;
    }
  }
  return output;
}

// 扁平化操作，将二维数组转换为一维数组，以便输入到全连接层
function flatten(input) {
  return input.flat();
}

// 全连接层前向传播，将输入数据与权重相乘并加上偏置，然后通过 sigmoid 函数进行激活
// 参数 input 是输入的一维数据，weights 是全连接层的权重矩阵，biases 是偏置向量
function fullyConnected(input, weights, biases) {
  const output = [];
  // 遍历输出的每个元素
  for (let i = 0; i < weights.length; i++) {
    let sum = biases[i];
    // 遍历输入的每个元素
    for (let j = 0; j < input.length; j++) {
      // 计算输入元素与对应权重的乘积并累加到总和中
      sum += input[j] * weights[i][j];
    }
    // 对总和应用 sigmoid 激活函数并存入输出数组
    output[i] = sigmoid(sum);
  }
  return output;
}

// 交叉熵损失函数，用于计算预测结果与真实标签之间的损失
// 参数 output 是模型的预测输出，label 是真实标签
function crossEntropyLoss(output, label) {
  // 创建一个长度为输出长度的数组，将真实标签对应的位置置为 1，其余位置置为 0
  const target = new Array(output.length).fill(0);
  target[label] = 1;
  let loss = 0;
  // 遍历输出的每个元素
  for (let i = 0; i < output.length; i++) {
    // 计算交叉熵损失并累加
    loss += -target[i] * Math.log(output[i]);
  }
  return loss;
}

// 训练函数，实现整个卷积神经网络的训练过程
// 参数 images 是输入的图像数据数组，labels 是对应的标签数组，numEpochs 是训练轮数，learningRate 是学习率
function train(images, labels, numEpochs = 10, learningRate = 0.01) {
  // 初始化卷积核
  const kernel = [
    [1, 1, 1],
    [1, 1, 1],
    [1, 1, 1]
  ];

  // 初始化全连接层的参数
  const numClasses = 10; // 分类的类别数
  // 计算扁平化后的输入大小
  const flattenedSize = Math.floor((28 - kernel.length + 1) / 2) * Math.floor((28 - kernel.length + 1) / 2);
  const weights = []; // 全连接层的权重矩阵
  const biases = []; // 全连接层的偏置向量
  // 初始化权重和偏置
  for (let i = 0; i < numClasses; i++) {
    weights[i] = [];
    for (let j = 0; j < flattenedSize; j++) {
      // 随机初始化权重，范围在 (-0.5, 0.5) 之间
      weights[i][j] = Math.random() - 0.5;
    }
    // 随机初始化偏置，范围在 (-0.5, 0.5) 之间
    biases[i] = Math.random() - 0.5;
  }

  // 开始训练，进行指定轮数的迭代
  for (let epoch = 0; epoch < numEpochs; epoch++) {
    let totalLoss = 0;
    // 遍历每一个训练样本
    for (let i = 0; i < images.length; i++) {
      // 将一维的图像数据转换为二维数组，并进行归一化处理
      const image = convert1Dto2D(images[i].data.map(pixel => pixel / 255), images[i].width);
      const label = labels[i];

      // 前向传播过程
      // 进行卷积操作
      const convOutput = convolve(image, kernel);
      // 进行最大池化操作
      const pooledOutput = maxPooling(convOutput);
      // 将池化后的二维数据扁平化
      const flattenedOutput = flatten(pooledOutput);
      // 通过全连接层得到最终的预测输出
      const fcOutput = fullyConnected(flattenedOutput, weights, biases);

      // 计算当前样本的损失
      const loss = crossEntropyLoss(fcOutput, label);
      totalLoss += loss;

      // 反向传播过程（简化实现，仅更新全连接层）
      // 创建一个长度为类别数的数组，将真实标签对应的位置置为 1，其余位置置为 0
      const target = new Array(numClasses).fill(0);
      target[label] = 1;
      const outputError = [];
      // 计算输出层的误差
      for (let j = 0; j < numClasses; j++) {
        outputError[j] = (fcOutput[j] - target[j]) * sigmoidDerivative(fcOutput[j]);
      }

      // 更新全连接层的权重和偏置
      for (let j = 0; j < numClasses; j++) {
        // 更新偏置
        biases[j] -= learningRate * outputError[j];
        for (let k = 0; k < flattenedOutput.length; k++) {
          // 更新权重
          weights[j][k] -= learningRate * outputError[j] * flattenedOutput[k];
        }
      }
    }

    // 输出当前轮数的平均损失
    console.log(`Epoch ${epoch + 1}: Loss = ${totalLoss / images.length}`);
  }

  // 返回训练好的全连接层的权重和偏置
  return { weights, biases };
}

// 保存模型到文件
function saveModel(model, filePath) {
    const modelData = JSON.stringify(model);
    fs.writeFileSync(filePath, modelData);
    console.log(`Model saved to ${filePath}`);
}

// 主函数
async function main() {
  const labels = await getUbyteFileData('./MINIST/train-labels-idx1-ubyte.gz');
  const images = await getUbyteFileData('./MINIST/train-images-idx3-ubyte.gz');
  const numEpochs = 10; // 训练轮数
  const learningRate = 0.01; // 学习率
  // 调用训练函数进行训练
  const trainedModel = train(images, labels, numEpochs, learningRate);
  console.log('Training completed. Trained model:', trainedModel);

  // 保存模型
  const modelFilePath = 'model.json';
  saveModel(trainedModel, modelFilePath);
  console.log('训练完成');
}
main();
